import os.path

from lib.trainer import Trainer
from lib.config import *
from lib.dataset import SpeakerDataset
from lib.dataset_spliter import DatasetSpliter

# def train_mfcc(model_type: str = 'resnet34', suffix: str = ''):
#     mfcc_feature_path = os.path.join(featureRoot, 'mfcc')
#     mfcc_model_path = os.path.join(modelRoot, 'mfcc', f'{model_type}_{suffix}')
#     if not os.path.exists(mfcc_model_path):
#         os.makedirs(mfcc_model_path)

#     trainer = Trainer(model_type, mfcc_feature_path, mfcc_model_path)
#     trainer.train()


fbank_feature_path = os.path.join(featureRoot, 'fbank')
dataset = SpeakerDataset(fbank_feature_path)
train_loader, test_loader = DatasetSpliter(dataset, trainPram['batchSize']).split()


def train_fbank(model_type, suffix: str = '', resume: bool = False):
    fbank_model_path = os.path.join(modelRoot, 'fbank', f'{model_type}{"_" if suffix != "" else ""}{suffix}')
    if not os.path.exists(fbank_model_path):
        os.makedirs(fbank_model_path)

    trainer = Trainer(model_type, fbank_model_path, train_loader, test_loader, resume)
    trainer.train()


# train_mfcc()
# train_fbank('resnet18')
# train_fbank('cnn18')
# train_fbank('cnn34')
# train_fbank('cnn50')
train_fbank('resnet34_custom', resume=True)
train_fbank('resnet50_custom', resume=True)
